package com.jobcn.spider;

import com.alibaba.fastjson.JSON;
import com.google.common.hash.BloomFilter;
import com.google.common.hash.Funnels;
import org.apache.commons.codec.digest.DigestUtils;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.scheduler.DuplicateRemovedScheduler;
import us.codecraft.webmagic.scheduler.MonitorableScheduler;
import us.codecraft.webmagic.scheduler.component.DuplicateRemover;

import java.nio.charset.Charset;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Url队列，使用BloomFilter去重
 * @author winson
 */
public class Scheduler58 extends DuplicateRemovedScheduler implements MonitorableScheduler, DuplicateRemover {

    protected JedisPool pool;

    private static final String QUEUE_PREFIX = "queue_";

    private static final String SET_PREFIX = "set_";

    private static final String ITEM_PREFIX = "item_";

    //布隆过滤器
    private long expectedInsertions;
    private double fpp;
    private BloomFilter<CharSequence> bloomFilter;
    private AtomicInteger counter;

    public Scheduler58(String host) {
        this(new JedisPool(new JedisPoolConfig(), host));
        this.expectedInsertions = 3500000L;
        this.fpp = 0.000001;
        bloomFilter = createBloomFilter();
        logger.info("布隆过滤器初始化完成...期望值:{},误判率:{}", expectedInsertions, fpp);
    }

    public Scheduler58(JedisPool pool) {
        this.pool = pool;
        setDuplicateRemover(this);
    }

    /**
     * 创建不布隆过滤器
     * @return
     */
    private BloomFilter<CharSequence> createBloomFilter() {
        counter = new AtomicInteger(0);
        return BloomFilter.create(
                Funnels.stringFunnel(Charset.defaultCharset()), expectedInsertions, fpp);
    }

    @Override
    public void resetDuplicateCheck(Task task) {
        Jedis jedis = pool.getResource();
        try {
            jedis.del(getSetKey(task));
        } finally {
            pool.returnResource(jedis);
        }
    }


    @Override
    public boolean isDuplicate(Request request, Task task) {
        String url = request.getUrl();
        boolean isDuplicate = mightContain(url);
        if (!isDuplicate) {
            counter.incrementAndGet();
        }
        return isDuplicate;
    }

    /**
     * 判断是否可能包含某个url
     *
     * @param str
     * @return true url可能已经存在
     *         false url一定不存在
     */
    public boolean mightContain(String str){
        boolean isDuplicate = false;
        synchronized (this) {
            isDuplicate = bloomFilter.mightContain(str);
            if (!isDuplicate) {
                bloomFilter.put(str);
            }
        }
        return isDuplicate;
    }


    @Override
    protected void pushWhenNoDuplicate(Request request, Task task) {
        Jedis jedis = pool.getResource();
        try {
            jedis.rpush(getQueueKey(task), request.getUrl());
            if (request.getExtras() != null) {
                String field = DigestUtils.shaHex(request.getUrl());
                String value = JSON.toJSONString(request);
                jedis.hset((ITEM_PREFIX + "www.58.com"), field, value);
            }
        } finally {
            pool.returnResource(jedis);
        }
    }

    @Override
    public synchronized Request poll(Task task) {
        Jedis jedis = pool.getResource();
        try {
            String url = jedis.lpop(getQueueKey(task));
            if (url == null) {
                return null;
            }
            String key = ITEM_PREFIX + "www.58.com";
            String field = DigestUtils.shaHex(url);
            byte[] bytes = jedis.hget(key.getBytes(), field.getBytes());
            if (bytes != null) {
                Request o = JSON.parseObject(new String(bytes), Request.class);
                //取出来就删掉
                jedis.hdel(key.getBytes(), field.getBytes());
                return o;
            }
                Request request = new Request(url);
            return request;
        } finally {
            pool.returnResource(jedis);
        }
    }

    protected String getSetKey(Task task) {return SET_PREFIX + "www.58.com";}

    protected String getQueueKey(Task task) {
        return QUEUE_PREFIX + "www.58.com";
    }

    protected String getItemKey(Task task)
    {
        return ITEM_PREFIX + "www.58.com";
    }

    @Override
    public int getLeftRequestsCount(Task task) {
        Jedis jedis = pool.getResource();
        try {
            Long size = jedis.llen(getQueueKey(task));
            return size.intValue();
        } finally {
            pool.returnResource(jedis);
        }
    }

    @Override
    public int getTotalRequestsCount(Task task) {
        Jedis jedis = pool.getResource();
        try {
            Long size = jedis.scard(getSetKey(task));
            return size.intValue();
        } finally {
            pool.returnResource(jedis);
        }
    }
}
